import { vec3 } from 'gl-matrix';
import arrowShader from './shaders/arrow_shader.wgsl?raw';
import spotlightConeShader from './shaders/spotlight_cone_shader.wgsl?raw';
import { createGPUBuffer, shaderModuleFromCode } from './utils';
import { Renderer } from './renderer';

export interface ISpotLight {
  pos: Float32Array;
  dir: Float32Array;
  color: Float32Array;
}

export class SpotLight {
  renderer: Renderer;
  lights: Map<number, ISpotLight>;
  currentId: number;

  shaderModule: GPUShaderModule;
  shaderModule2: GPUShaderModule;

  instanceDataBuffer?: GPUBuffer;
  instanceDataBuffer2?: GPUBuffer;
  positionBuffer1: GPUBuffer;
  positionBuffer2: GPUBuffer;
  uniformBindGroupLayout: GPUBindGroupLayout;
  pipeline: GPURenderPipeline;
  pipeline2: GPURenderPipeline;
  uniformBindGroup: GPUBindGroup;

  constructor(renderer: Renderer) {
    this.renderer = renderer;
    this.lights = new Map();
    this.currentId = 0;

    const { device } = renderer;

    this.shaderModule = shaderModuleFromCode(device, arrowShader);
    this.shaderModule2 = shaderModuleFromCode(device, spotlightConeShader);

    const radius = 1.0;

    const positionBuffer1 = [];
    const positionBuffer2 = [];

    for (let i = 0; i < 10; ++i) {
      const angle = (i * 2.0 * Math.PI) / 10;
      const angleNext = ((i + 1) * 2.0 * Math.PI) / 10;
      const angleMiddle = ((i + 0.5) * 2.0 * Math.PI) / 10;

      const xS = Math.cos(angle) * radius;
      const yS = Math.sin(angle) * radius;

      const xSN = Math.cos(angleNext) * radius;
      const ySN = Math.sin(angleNext) * radius;

      const xSM = Math.cos(angleMiddle) * radius;
      const ySM = Math.sin(angleMiddle) * radius;
      const coneCos = 0.9;
      const s = Math.sqrt(1.0 - coneCos * coneCos) / coneCos;

      positionBuffer2.push(
        ...[0.0, 0.0, 0.0, xSM * s, 1.0 - radius * s * s * 4.0 * radius, ySM * s]
      );
      positionBuffer2.push(
        ...[xSN * s, 1.0, ySN * s, xSM * s, 1.0 - radius * s * s * 4.0 * radius, ySM * s]
      );
      positionBuffer2.push(
        ...[xS * s, 1.0, yS * s, xSM * s, 1.0 - radius * s * s * 4.0 * radius, ySM * s]
      );

      positionBuffer1.push(...[xSN * s, 1.0, ySN * s]);
      positionBuffer1.push(...[xS * s, 1.0, yS * s]);
      positionBuffer1.push(...[xS * s, 1.0, yS * s]);
      positionBuffer1.push(...[0.0, 0.0, 0.0]);
    }

    this.positionBuffer1 = createGPUBuffer(
      device,
      new Float32Array(positionBuffer1),
      GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    );
    this.positionBuffer2 = createGPUBuffer(
      device,
      new Float32Array(positionBuffer2),
      GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    );

    this.uniformBindGroupLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.VERTEX, buffer: {} },
        { binding: 1, visibility: GPUShaderStage.VERTEX, buffer: {} },
        { binding: 2, visibility: GPUShaderStage.VERTEX, buffer: {} },
        { binding: 3, visibility: GPUShaderStage.VERTEX, buffer: {} },
      ],
    });

    const layout = device.createPipelineLayout({ bindGroupLayouts: [this.uniformBindGroupLayout] });

    this.pipeline = device.createRenderPipeline({
      layout,
      vertex: {
        module: this.shaderModule2,
        entryPoint: 'vs_main',
        buffers: [
          {
            attributes: [
              {
                shaderLocation: 0, // @location(0)
                offset: 0,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 3, // sizeof(float) * 6
            stepMode: 'vertex',
          },
          {
            attributes: [
              {
                shaderLocation: 1, // @location(1)
                offset: 0,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 7, // sizeof(float) * 6
            stepMode: 'instance',
          },
          {
            attributes: [
              {
                shaderLocation: 2, // @location(1)
                offset: 3 * 4,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 7, // sizeof(float) * 6
            stepMode: 'instance',
          },
          {
            attributes: [
              {
                shaderLocation: 3, // @location(1)
                offset: 6 * 4,
                format: 'float32',
              },
            ],
            arrayStride: 4 * 7, // sizeof(float) * 6
            stepMode: 'instance',
          },
        ],
      },
      fragment: {
        module: this.shaderModule2,
        entryPoint: 'fs_main',
        targets: [{ format: 'bgra8unorm' }],
      },
      primitive: {
        topology: 'line-list',
      },
      depthStencil: {
        depthWriteEnabled: true,
        depthCompare: 'less',
        format: 'depth32float',
      },
    });

    this.pipeline2 = device.createRenderPipeline({
      layout,
      vertex: {
        module: this.shaderModule,
        entryPoint: 'vs_main',
        buffers: [
          {
            attributes: [
              {
                shaderLocation: 0, // @location(0)
                offset: 0,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 6, // sizeof(float) * 6
            stepMode: 'vertex',
          },
          {
            attributes: [
              {
                shaderLocation: 1, // @location(1)
                offset: 4 * 3,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 6, // sizeof(float) * 6
            stepMode: 'vertex',
          },
          {
            attributes: [
              {
                shaderLocation: 2, // @location(1)
                offset: 0,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 9, // sizeof(float) * 6
            stepMode: 'instance',
          },
          {
            attributes: [
              {
                shaderLocation: 3, // @location(1)
                offset: 3 * 4,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 9, // sizeof(float) * 6
            stepMode: 'instance',
          },
          {
            attributes: [
              {
                shaderLocation: 4, // @location(1)
                offset: 6 * 4,
                format: 'float32x3',
              },
            ],
            arrayStride: 4 * 9, // sizeof(float) * 6
            stepMode: 'instance',
          },
        ],
      },
      fragment: {
        module: this.shaderModule,
        entryPoint: 'fs_main',
        targets: [{ format: 'bgra8unorm' }],
      },
      primitive: {
        topology: 'triangle-list',
        frontFace: 'ccw',
        cullMode: 'none',
      },
      depthStencil: {
        depthWriteEnabled: true,
        depthCompare: 'less',
        format: 'depth32float',
      },
    });

    const {
      modelViewMatrixUniformBuffer,
      projectionMatrixUniformBuffer,
      normalMatrixUniformBuffer,
      viewDirectionUniformBuffer,
    } = renderer.camera;

    this.uniformBindGroup = device.createBindGroup({
      layout: this.uniformBindGroupLayout,
      entries: [
        { binding: 0, resource: { buffer: modelViewMatrixUniformBuffer } },
        { binding: 1, resource: { buffer: projectionMatrixUniformBuffer } },
        { binding: 2, resource: { buffer: normalMatrixUniformBuffer } },
        { binding: 3, resource: { buffer: viewDirectionUniformBuffer } },
      ],
    });
  }

  upsertSpotLight(
    id: number | undefined,
    pos: Float32Array,
    dir: Float32Array,
    color: Float32Array
  ) {
    let assignedId = this.currentId;

    if (id != undefined && this.lights.has(id)) {
      assignedId = id;
    } else {
      this.currentId++;
    }

    this.lights.set(assignedId, { pos, dir, color });

    return assignedId;
  }

  refreshBuffer(device: GPUDevice) {
    this.instanceDataBuffer?.destroy();
    this.instanceDataBuffer2?.destroy();

    const instanceData = [];
    const instanceData2 = [];

    for (const value of this.lights.values()) {
      instanceData.push(value.pos[0]);
      instanceData.push(value.pos[1]);
      instanceData.push(value.pos[2]);

      instanceData.push(value.dir[0]);
      instanceData.push(value.dir[1]);
      instanceData.push(value.dir[2]);

      instanceData.push(value.color[0]);
      instanceData.push(value.color[1]);
      instanceData.push(value.color[2]);

      instanceData2.push(value.pos[0]);
      instanceData2.push(value.pos[1]);
      instanceData2.push(value.pos[2]);

      instanceData2.push(value.dir[0]);
      instanceData2.push(value.dir[1]);
      instanceData2.push(value.dir[2]);

      instanceData2.push(vec3.length(value.pos));
    }

    this.instanceDataBuffer = createGPUBuffer(
      device,
      new Float32Array(instanceData),
      GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    );
    this.instanceDataBuffer2 = createGPUBuffer(
      device,
      new Float32Array(instanceData2),
      GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    );
  }

  render(encoder: GPURenderPassEncoder) {
    if (this.lights.size > 0 && this.instanceDataBuffer && this.instanceDataBuffer2) {
      encoder.setPipeline(this.pipeline);
      encoder.setBindGroup(0, this.uniformBindGroup);
      encoder.setVertexBuffer(0, this.positionBuffer1);
      encoder.setVertexBuffer(1, this.instanceDataBuffer2);
      encoder.setVertexBuffer(2, this.instanceDataBuffer2);
      encoder.setVertexBuffer(3, this.instanceDataBuffer2);
      encoder.draw(40, this.lights.size);

      encoder.setPipeline(this.pipeline2);
      encoder.setBindGroup(0, this.uniformBindGroup);
      encoder.setVertexBuffer(0, this.positionBuffer2);
      encoder.setVertexBuffer(1, this.positionBuffer2);
      encoder.setVertexBuffer(2, this.instanceDataBuffer);
      encoder.setVertexBuffer(3, this.instanceDataBuffer);
      encoder.setVertexBuffer(4, this.instanceDataBuffer);
      encoder.draw(30, this.lights.size);
    }
  }
}
